import os
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
import warnings
import glob
import scipy.io
import re

warnings.filterwarnings('ignore')

class ChannelDataset(Dataset):
    """
    信道特征数据集加载器
    处理包含多簇信道特征和探头权重标签的数据
    """
    # 添加静态变量用于缓存文件数据
    _file_cache = None
    _spatial_corr_cache = {}
    
    def __init__(self, root_path, flag='train', size=None, 
                 features='M', data_path='channel_data.npz', 
                 n_clusters=25, loader='channel', load_spatial_corr=True,
                 base_data=None):
        """
        初始化信道特征数据集
        
        Args:
            root_path (str): 数据集根目录（不再使用，保留参数兼容性）
            flag (str): 'train'/'val'/'test'/'all' 指定加载哪个部分的数据
            size (list): [seq_len, label_len, pred_len] 分别表示输入序列长度、标签长度和预测长度
            features (str): 特征类型，默认为'M'多变量
            data_path (str): 数据文件名（不再使用，保留参数兼容性）
            n_clusters (int): 每个时间点的簇数量，默认为25
            loader (str): 数据加载类型
            load_spatial_corr (bool): 是否加载空间相关性数据
            base_data (object): 基础数据对象，如果提供则直接使用而不重新加载
        """
        # 初始化大小信息
        if size == None:
            self.seq_len = 20  # 默认为20个时间点的序列
            self.label_len = 10
            self.pred_len = 10
        else:
            self.seq_len = size[0]
            self.label_len = size[1]
            self.pred_len = size[2]
            
        # 初始化信息
        assert flag in ['train', 'test', 'val', 'all']
        type_map = {'train': 0, 'val': 1, 'test': 2, 'all': 3}
        self.set_type = type_map[flag]
        self.flag = flag
        
        self.features = features
        self.root_path = root_path  # 保留但不使用
        self.data_path = data_path  # 保留但不使用
        self.n_clusters = n_clusters
        self.feature_dim = 5  # 每个簇有5个特征
        self.load_spatial_corr = load_spatial_corr
        
        # 加载数据（只调用一次）
        self._load_data(loader)
    
    def _load_data(self, loader):
        """
        加载数据逻辑
        
        Args:
            loader (str): 加载器类型
        """
        # 打印空间相关性数据加载状态
        print(f"空间相关性数据加载状态: {self.load_spatial_corr}")
        
        try:
            # 加载数据
            if loader == 'channel':
                # 检查是否已有缓存
                if ChannelDataset._file_cache is not None:
                    print("使用缓存的CSV数据文件，跳过文件IO")
                    # 使用已缓存的数据
                    df = ChannelDataset._file_cache
                    time_points = sorted(df['receiver_point'].unique())
                    n_time_points = len(time_points)
                    feature_columns = [
                        'horizontal_aoa',  # 水平到达角
                        'vertical_aoa',    # 垂直到达角
                        'received_power',  # 接收功率(dB)
                        'ASA',             # 水平扩展角
                        'ZSA'              # 垂直扩展角
                    ]
                    max_clusters = int(df['path_number'].max())
                else:
                    # 首先检查CSV文件是否存在
                    csv_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'merged_output.csv')
                    if not os.path.exists(csv_path):
                        raise FileNotFoundError(f"找不到CSV数据文件: {csv_path}")
                    
                    print(f"第一次加载CSV数据文件: {csv_path}")
                    
                    # 使用pandas加载CSV数据
                    import pandas as pd
                    df = pd.read_csv(csv_path)
                    
                    # 缓存数据框
                    ChannelDataset._file_cache = df
                    
                    # 提取所有唯一的时间点ID
                    time_points = sorted(df['receiver_point'].unique())
                    n_time_points = len(time_points)
                    print(f"CSV中的总时间点数: {n_time_points}")
                    
                    if n_time_points < 1:
                        raise ValueError(f"CSV文件中没有有效的时间点数据")
                    
                    # 特征列映射
                    feature_columns = [
                        'horizontal_aoa',  # 水平到达角
                        'vertical_aoa',    # 垂直到达角
                        'received_power',  # 接收功率(dB)
                        'ASA',             # 水平扩展角
                        'ZSA'              # 垂直扩展角
                    ]
                    
                    # 输出一些统计信息帮助调试
                    cluster_counts = df.groupby('receiver_point')['path_number'].max()
                    avg_clusters = cluster_counts.mean()
                    print(f"平均每个时间点的簇数量: {avg_clusters:.2f}")
                    
                    # 确定最大的簇数量
                    max_clusters = int(df['path_number'].max())
                    print(f"最大簇数量: {max_clusters}")
                
                # 计算窗口大小和窗口数量
                window_size = 20  # 默认窗口大小
                # 计算可以生成的完整窗口数量
                n_windows = n_time_points // window_size
                if n_windows < 1:
                    raise ValueError(f"时间点数量 ({n_time_points}) 不足以生成大小为 {window_size} 的窗口")
                
                print(f"使用窗口大小: {window_size}, 可生成窗口数量: {n_windows}")
                
                # 创建一个结构化的数据数组来存储所有特征
                # 形状: [n_windows, window_size, n_clusters, 5]
                # 其中5表示5个特征维度
                n_clusters = min(max_clusters, self.n_clusters)  # 使用参数中指定的簇数量，不超过最大值
                print(f"使用簇数量: {n_clusters}")
                
                # 初始化数据数组
                channel_features = np.zeros((n_windows, window_size, n_clusters, 5), dtype=np.float32)
                # 初始化时间点ID数组
                time_ids = np.zeros((n_windows, window_size), dtype=np.int32)
                
                # 如果需要加载空间相关性数据，提前准备
                spatial_corr_dict = {}
                if self.load_spatial_corr:
                    # 如果已有缓存，直接使用
                    if ChannelDataset._spatial_corr_cache:
                        print("使用缓存的空间相关性数据，跳过文件IO")
                        spatial_corr_dict = ChannelDataset._spatial_corr_cache
                    else:
                        # 修正空间相关性数据目录路径
                        spatial_corr_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'dy')
                        print(f"尝试加载空间相关性数据，路径: {spatial_corr_dir}")
                        
                        if os.path.exists(spatial_corr_dir):
                            print(f"加载空间相关性数据从: {spatial_corr_dir}")
                            # 获取所有空间相关性文件
                            spatial_files = glob.glob(os.path.join(spatial_corr_dir, 'spatial_correlation_*.mat'))
                            print(f"找到{len(spatial_files)}个空间相关性数据文件")
                            
                            # 显示前几个文件名以便调试
                            if spatial_files:
                                file_examples = [os.path.basename(f) for f in spatial_files[:5]]
                                print(f"文件样本(前5个): {file_examples}")
                            
                            # 预加载所有空间相关性数据到字典中
                            for file_path in spatial_files:
                                # 从文件名中提取时间点ID
                                match = re.search(r'spatial_correlation_(\d+)\.mat', os.path.basename(file_path))
                                if match:
                                    time_id = int(match.group(1))
                                    try:
                                        # 加载.mat文件中的空间相关性数据
                                        mat_data = scipy.io.loadmat(file_path)
                                        if "p" in mat_data:
                                            # 直接使用原始的p值，不做任何数值转换
                                            spatial_corr = mat_data['p']  # 形状应为(1, 64)
                                            # 储存原始数据
                                            spatial_corr_dict[time_id] = spatial_corr.flatten()
                                        else:
                                            print(f"警告: 文件 {file_path} 中没有找到 'p' 键")
                                    except Exception as e:
                                        print(f"加载空间相关性文件 {file_path} 失败: {e}")
                            
                            print(f"成功预加载 {len(spatial_corr_dict)} 个空间相关性数据点")
                            if spatial_corr_dict:
                                loaded_ids = list(spatial_corr_dict.keys())
                                print(f"空间相关性数据时间点ID范围: {min(loaded_ids)} - {max(loaded_ids)}")
                                
                                # 缓存空间相关性数据
                                ChannelDataset._spatial_corr_cache = spatial_corr_dict
                        else:
                            print(f"空间相关性数据目录不存在: {spatial_corr_dir}")
                            self.load_spatial_corr = False
                
                # 初始化空间相关性数据数组（如果需要）
                if self.load_spatial_corr and spatial_corr_dict:
                    # 假设每个空间相关性数据点是64维复数向量
                    spatial_corr_real = np.zeros((n_windows, window_size, 64), dtype=np.float32)
                    spatial_corr_imag = np.zeros((n_windows, window_size, 64), dtype=np.float32)
                    missing_corr_count = 0
                    total_corr_count = 0
                
                # 对于每个窗口
                for w_idx in range(n_windows):
                    # 对于窗口中的每个时间点
                    for t_idx in range(window_size):
                        # 计算当前时间点在总序列中的索引
                        time_idx = w_idx * window_size + t_idx
                        
                        # 安全检查，确保索引在范围内
                        if time_idx >= len(time_points):
                            break
                            
                        # 获取当前时间点ID
                        current_time_id = time_points[time_idx]
                        time_ids[w_idx, t_idx] = current_time_id
                        
                        # 获取当前时间点的所有数据
                        time_data = df[df['receiver_point'] == current_time_id]
                        
                        # 按路径编号排序
                        sorted_paths = time_data.sort_values('path_number')
                        
                        # 取前n_clusters个簇
                        for c_idx in range(min(len(sorted_paths), n_clusters)):
                            if c_idx < len(sorted_paths):
                                cluster = sorted_paths.iloc[c_idx]
                                
                                # 提取特征
                                for f_idx, feature in enumerate(feature_columns):
                                    if feature in cluster:
                                        channel_features[w_idx, t_idx, c_idx, f_idx] = cluster[feature]
                        
                        # 同时加载对应的空间相关性数据（如果有）
                        if self.load_spatial_corr and spatial_corr_dict:
                            total_corr_count += 1
                            if current_time_id in spatial_corr_dict:
                                # 获取对应的空间相关性数据
                                corr_data = spatial_corr_dict[current_time_id]
                                
                                # 分离实部和虚部
                                spatial_corr_real[w_idx, t_idx] = np.real(corr_data)
                                spatial_corr_imag[w_idx, t_idx] = np.imag(corr_data)
                            else:
                                missing_corr_count += 1
                                # 如果找不到对应的数据，用零填充
                                spatial_corr_real[w_idx, t_idx] = 0
                                spatial_corr_imag[w_idx, t_idx] = 0
                
                # 输出数据形状
                print(f"成功生成数据数组，形状: {channel_features.shape}")
                print(f"时间点ID数组形状: {time_ids.shape}")
                
                # 输出时间点ID的范围和样本
                time_id_min = time_ids.min()
                time_id_max = time_ids.max()
                print(f"时间点ID范围: {time_id_min} - {time_id_max}")
                time_id_examples = np.unique(time_ids.flatten())[:10]
                print(f"时间点ID样本(前10个): {time_id_examples}")
                
                # 如果加载了空间相关性数据，输出统计信息
                if self.load_spatial_corr and spatial_corr_dict:
                    print(f"空间相关性数据统计: 总时间点 {total_corr_count}, 缺失时间点 {missing_corr_count}, 缺失率 {missing_corr_count/total_corr_count*100:.2f}%")
                    # 保存空间相关性字典，以便后续查询
                    self.spatial_corr_dict = spatial_corr_dict
                
                # 处理数据，区分训练/验证/测试集
                # 使用窗口作为样本单位，而不是时间点
                num_samples = n_windows
                
                # 创建窗口索引数组并打乱顺序
                window_indices = np.arange(num_samples)
                np.random.seed(42)  # 设置随机种子以确保结果可重现
                np.random.shuffle(window_indices)
                
                # 按打乱后的顺序重新排列数据
                channel_features = channel_features[window_indices]
                time_ids = time_ids[window_indices]
                if self.load_spatial_corr and 'spatial_corr_real' in locals() and 'spatial_corr_imag' in locals():
                    spatial_corr_real = spatial_corr_real[window_indices]
                    spatial_corr_imag = spatial_corr_imag[window_indices]
                
                print(f"已打乱窗口顺序，确保训练集包含整个时间范围的数据")
                
                # 现在按比例划分数据集
                border1s = [0, int(num_samples*0.7), int(num_samples*0.7), 0]  # 验证集和测试集使用相同数据（后30%）
                border2s = [int(num_samples*0.7), num_samples, num_samples, num_samples]
                border1 = border1s[self.set_type]
                border2 = border2s[self.set_type]
                
                # 划分数据集
                self.data_x = channel_features[border1:border2]  # [n_samples, window_size, n_clusters, 5]
                self.time_ids = time_ids[border1:border2]  # [n_samples, window_size]
                
                # 如果有空间相关性数据，也进行划分
                if self.load_spatial_corr and 'spatial_corr_real' in locals() and 'spatial_corr_imag' in locals():
                    self.spatial_corr_real = spatial_corr_real[border1:border2]  # [n_samples, window_size, 64]
                    self.spatial_corr_imag = spatial_corr_imag[border1:border2]  # [n_samples, window_size, 64]
                
                print(f"划分后数据形状 - data_x: {self.data_x.shape}")
                print(f"划分后时间点ID形状: {self.time_ids.shape}")
                print(f"划分后时间点ID范围: {self.time_ids.min()} - {self.time_ids.max()}")
                
                # 输出训练集中时间点ID的分布情况
                if self.set_type == 0:  # 训练集
                    unique_time_ids = np.unique(self.time_ids.flatten())
                    print(f"训练集包含的时间点ID数量: {len(unique_time_ids)}")
                    print(f"训练集时间点ID范围: {unique_time_ids.min()} - {unique_time_ids.max()}")
                    # 显示时间点ID的分布
                    time_id_samples = unique_time_ids[::len(unique_time_ids)//10] if len(unique_time_ids) > 10 else unique_time_ids
                    print(f"训练集时间点ID分布样本: {time_id_samples}")
            else:
                raise ValueError(f"不支持的数据加载方式: {loader}")
            
        except Exception as e:
            # 发生错误时，打印错误并向上抛出
            print(f"加载数据失败: {e}")
            raise RuntimeError(f"数据加载失败: {str(e)}") from e
        
        # 确保数据类型
        self.data_x = self.data_x.astype(np.float32)
        self.time_stamps = np.expand_dims(np.arange(self.seq_len), axis=0)
        self.time_stamps = np.repeat(self.time_stamps, len(self.data_x), axis=0)  # [n_samples, seq_len]
    
    def __getitem__(self, index):
        """
        获取单个样本
        
        Args:
            index (int): 窗口索引
            
        Returns:
            tuple: (x_dict, y_dict, x_mark, y_mark) 包含输入、目标和对应的时间特征
        """
        # 直接获取对应窗口的数据
        window_x = self.data_x[index]  # [window_size, n_clusters, 5]
        window_time = self.time_stamps[index]  # [seq_len]
        
        # 获取当前窗口的时间点ID
        window_time_ids = self.time_ids[index]  # [window_size]
        
        # 使用全部序列作为输入
        seq_x = window_x[:self.seq_len]  # [seq_len, n_clusters, 5]
        seq_time_ids = window_time_ids[:self.seq_len]  # [seq_len]
        
        # 时间戳，使用全序列的时间标记
        seq_x_mark = window_time[:self.seq_len]
        
        # 为了适配模型，重新整形特征，将簇和特征维度展平
        seq_x_flat = seq_x.reshape(seq_x.shape[0], -1)  # [seq_len, n_clusters*5]
        
        # 处理空间相关性数据
        if self.load_spatial_corr and hasattr(self, 'spatial_corr_real') and hasattr(self, 'spatial_corr_imag'):
            # 直接使用预加载的空间相关性数据
            spatial_corrs_real = self.spatial_corr_real[index][:self.seq_len]  # [seq_len, 64]
            spatial_corrs_imag = self.spatial_corr_imag[index][:self.seq_len]  # [seq_len, 64]
        else:
            # 如果没有预加载的空间相关性数据，使用零数组
            spatial_corrs_real = np.zeros((self.seq_len, 64), dtype=np.float32)
            spatial_corrs_imag = np.zeros((self.seq_len, 64), dtype=np.float32)
        
        # 处理为Pytorch张量
        x_dict = {
            'data': torch.FloatTensor(seq_x_flat),
            'orig_data': torch.FloatTensor(seq_x),  # 保留原始形状以便需要时使用
            'time_ids': torch.LongTensor(seq_time_ids),  # 添加时间点ID
            'spatial_corr_real': torch.FloatTensor(spatial_corrs_real),  # 空间相关性实部
            'spatial_corr_imag': torch.FloatTensor(spatial_corrs_imag)   # 空间相关性虚部
        }
        
        return x_dict, torch.FloatTensor(seq_x_mark).unsqueeze(-1)
    
    def __len__(self):
        """
        返回数据集长度
        
        Returns:
            int: 数据集中样本窗口数量
        """
        return len(self.data_x) 